{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore\n",
    "from mindspore import nn\n",
    "from mindspore.dataset import vision, transforms\n",
    "from mindspore.dataset import MnistDataset\n",
    "from mindspore.train import Model, CheckpointConfig, ModelCheckpoint, LossMonitor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "current dir /home/wm\n",
      "change dir\n",
      "/home/wm/statebear/jupyter/mindspore\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'/home/wm/statebear/jupyter/mindspore'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    " # 切换目录\n",
    "c_dir=%pwd\n",
    "print(\"current dir\",c_dir)\n",
    "if ('mindspore' not in c_dir):\n",
    "    print(\"change dir\")\n",
    "    %cd statebear/jupyter/mindspore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/MNIST_Data.zip (10.3 MB)\n",
      "\n",
      "file_sizes: 100%|██████████████████████████| 10.8M/10.8M [00:00<00:00, 27.2MB/s]\n",
      "Extracting zip file...\n",
      "Successfully downloaded / unzipped to ./\n"
     ]
    }
   ],
   "source": [
    "# Download data from open datasets\n",
    "from download import download\n",
    "\n",
    "url = \"https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/\" \\\n",
    "      \"notebook/datasets/MNIST_Data.zip\"\n",
    "path = download(url, \"./\", kind=\"zip\", replace=True)\n",
    "\n",
    "\n",
    "def datapipe(path, batch_size):\n",
    "    image_transforms = [\n",
    "        vision.Rescale(1.0 / 255.0, 0),\n",
    "        vision.Normalize(mean=(0.1307,), std=(0.3081,)),\n",
    "        vision.HWC2CHW()\n",
    "    ]\n",
    "    label_transform = transforms.TypeCast(mindspore.int32)\n",
    "\n",
    "    dataset = MnistDataset(path)\n",
    "    dataset = dataset.map(image_transforms, 'image')\n",
    "    dataset = dataset.map(label_transform, 'label')\n",
    "    dataset = dataset.batch(batch_size)\n",
    "    return dataset\n",
    "\n",
    "train_dataset = datapipe('MNIST_Data/train', 64)\n",
    "test_dataset = datapipe('MNIST_Data/test', 64)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define model\n",
    "class Network(nn.Cell):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.dense_relu_sequential = nn.SequentialCell(\n",
    "            nn.Dense(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Dense(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Dense(512, 10)\n",
    "        )\n",
    "\n",
    "    def construct(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.dense_relu_sequential(x)\n",
    "        return logits\n",
    "\n",
    "model = Network()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate loss function and optimizer\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = nn.SGD(model.trainable_params(), 1e-2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_per_epoch = train_dataset.get_dataset_size()\n",
    "config = CheckpointConfig(save_checkpoint_steps=steps_per_epoch)\n",
    "\n",
    "ckpt_callback = ModelCheckpoint(prefix=\"mnist\", directory=\"./checkpoint\", config=config)\n",
    "loss_callback = LossMonitor(steps_per_epoch)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 step: 938, loss is 0.537405252456665\n",
      "Eval result: epoch 1, metrics: {'accuracy': 0.9034}\n",
      "epoch: 2 step: 938, loss is 0.16516736149787903\n",
      "Eval result: epoch 2, metrics: {'accuracy': 0.9291}\n",
      "epoch: 3 step: 938, loss is 0.17213881015777588\n",
      "Eval result: epoch 3, metrics: {'accuracy': 0.9372}\n",
      "epoch: 4 step: 938, loss is 0.2571163773536682\n",
      "Eval result: epoch 4, metrics: {'accuracy': 0.9421}\n",
      "epoch: 5 step: 938, loss is 0.06666439026594162\n",
      "Eval result: epoch 5, metrics: {'accuracy': 0.9535}\n",
      "epoch: 6 step: 938, loss is 0.1474754512310028\n",
      "Eval result: epoch 6, metrics: {'accuracy': 0.9563}\n",
      "epoch: 7 step: 938, loss is 0.3110397458076477\n",
      "Eval result: epoch 7, metrics: {'accuracy': 0.9608}\n",
      "epoch: 8 step: 938, loss is 0.06483744084835052\n",
      "Eval result: epoch 8, metrics: {'accuracy': 0.9642}\n",
      "epoch: 9 step: 938, loss is 0.03409641981124878\n",
      "Eval result: epoch 9, metrics: {'accuracy': 0.9658}\n",
      "epoch: 10 step: 938, loss is 0.035747140645980835\n",
      "Eval result: epoch 10, metrics: {'accuracy': 0.9683}\n"
     ]
    }
   ],
   "source": [
    "trainer = Model(model, loss_fn=loss_fn, optimizer=optimizer, metrics={'accuracy'})\n",
    "\n",
    "trainer.fit(10, train_dataset, test_dataset, callbacks=[ckpt_callback, loss_callback])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
